from langgraph.graph import END, StateGraph, START
from Agents.DialogState import State
from tools.get_users import get_user_info
from langchain_core.runnables import RunnableConfig
from Agents.PrimaryAgent import primary_assistant
from Agents.ProductAgent import product_assistant
from Agents.RouteTools import create_entry_node
from tools.utility import create_tool_node_with_fallback
from Agents.PrimaryAgent import primary_tools
from Agents.RouteTools import route_of_product_assistant
from langgraph.checkpoint.memory import MemorySaver
from Agents.RouteTools import route_to_workflow,route_primary_assistant
from Agents.AfterSaleAgent import after_sales_assistant
from Agents.ProductAgent import product_safe_tools,product_sensitive_tools
from Agents.AfterSaleAgent import after_sales_safe_tools,after_sales_sensitive_tools
from Agents.RouteTools import route_of_after_sales_assistant
from Agents.OrdersAgent import orders_assistant
from Agents.OrdersAgent import orders_safe_tools,orders_sensitive_tools
from Agents.RouteTools import route_of_orders_assistant

from Agents.RouteTools import pop_dialog_state
builder = StateGraph(State)


def fetch_user_info(state: State,config: RunnableConfig):
    configuration = config.get("configurable", {})
    customer_id = configuration.get("customer_id", None)
    return {"user_info": get_user_info.invoke({"customer_id": customer_id})}


builder.add_node("fetch_user_info", fetch_user_info)
builder.add_node("primary_assistant", primary_assistant)
builder.add_node("primary_assistant_tools", create_tool_node_with_fallback(primary_tools))

####增加开始节点的路由，路由到主代理和各个专业代理
builder.add_edge(START, "fetch_user_info")

builder.add_conditional_edges(
    "fetch_user_info",
    route_to_workflow,
)
####增加主代理的路由，路由到各个专业代理
builder.add_conditional_edges("primary_assistant", route_primary_assistant,
                    ["enter_product_assistant",
                    "enter_after_sales_assistant",
                    "enter_orders_assistant",
                    "primary_assistant_tools",
                    END])

builder.add_edge("primary_assistant_tools", "primary_assistant")

###构建leave_node图

builder.add_node("leave_node", pop_dialog_state)
builder.add_edge("leave_node", "primary_assistant")


####构建product_assistant图

builder.add_node("enter_product_assistant",
                     create_entry_node("Product Assistant", "product_assistant"))
builder.add_node("product_assistant", product_assistant)
builder.add_node("product_safe_tools", 
            create_tool_node_with_fallback(product_safe_tools))
builder.add_node("product_sensitive_tools", 
            create_tool_node_with_fallback(product_sensitive_tools))


builder.add_edge("enter_product_assistant", "product_assistant")
builder.add_conditional_edges(
    "product_assistant", route_of_product_assistant,
                    ["leave_node","product_safe_tools","product_sensitive_tools",END])

builder.add_edge("product_safe_tools", "product_assistant")
builder.add_edge("product_sensitive_tools", "product_assistant")


####构建after_sales_assistant图

builder.add_node("enter_after_sales_assistant",
                     create_entry_node("After Sales Assistant", "after_sales_assistant"))
builder.add_node("after_sales_assistant", after_sales_assistant)
builder.add_node("after_sales_safe_tools", 
        create_tool_node_with_fallback(after_sales_safe_tools))
builder.add_node("after_sales_sensitive_tools", 
        create_tool_node_with_fallback(after_sales_sensitive_tools))

builder.add_edge("enter_after_sales_assistant", "after_sales_assistant")
builder.add_conditional_edges(
    "after_sales_assistant", route_of_after_sales_assistant,
                    ["leave_node","after_sales_safe_tools","after_sales_sensitive_tools",END])

builder.add_edge("after_sales_safe_tools", "after_sales_assistant")
builder.add_edge("after_sales_sensitive_tools", "after_sales_assistant")
####构建orders_assistant图

builder.add_node("enter_orders_assistant",
                     create_entry_node("Orders Assistant", "orders_assistant"))
builder.add_node("orders_assistant", orders_assistant)
builder.add_node("orders_safe_tools", create_tool_node_with_fallback(orders_safe_tools))
builder.add_node("orders_sensitive_tools", 
                create_tool_node_with_fallback(orders_sensitive_tools))

builder.add_edge("enter_orders_assistant", "orders_assistant")
builder.add_conditional_edges(
    "orders_assistant", route_of_orders_assistant,
                    ["leave_node","orders_safe_tools","orders_sensitive_tools",END])

builder.add_edge("orders_safe_tools", "orders_assistant")
builder.add_edge("orders_sensitive_tools", "orders_assistant")

memory = MemorySaver()
graph = builder.compile(
    checkpointer=memory,
    interrupt_before=[
        "product_sensitive_tools",
        "after_sales_sensitive_tools",
        "orders_sensitive_tools",
    ],

)





try:
    graph.get_graph().draw_mermaid_png(output_file_path="main_graph.png")
except Exception:
    # This requires some extra dependencies and is optional
    pass


































